#include "mem/pmem.h"
#include "lib/print.h"
#include "lib/lock.h"
#include "lib/str.h"
#include "memlayout.h"
#include "common.h"
#include "mem/vmem.h"

// // 物理页节点
// typedef struct page_node {
//     struct page_node* next;
//     uint32 refNum; // 引用计数
// } page_node_t;

// 许多物理页构成一个可分配的区域
typedef struct alloc_region {
    uint64 begin;          // 起始物理地址
    uint64 end;            // 终止物理地址
    spinlock_t lk;         // 自旋锁(保护下面两个变量)
    uint32 allocable;      // 可分配页面数    
    page_node_t list_head; // 可分配链的链头节点
} alloc_region_t;

// 内核和用户可分配的物理页分开
static alloc_region_t kern_region, user_region;

#define KERN_PAGES 1024 // 内核可分配空间占1024个pages

void init_range(void *pa_start, void *pa_end,bool in_kernel);

// 物理内存初始化
void pmem_init(void)
{
    kern_region.begin = (uint64)ALLOC_BEGIN; // 物理内存分配的起始地址
    kern_region.end = (uint64)(ALLOC_BEGIN + (KERN_PAGES * PGSIZE)-1); 

    kern_region.allocable = (kern_region.end - kern_region.begin) / PGSIZE +1; // 计算可分配页面数
    kern_region.list_head.next = NULL; // 初始化链表头为 NULL
    spinlock_init(&kern_region.lk, "kernel_region_lock"); // 初始化自旋锁

    // 初始化用户区域
    user_region.begin = (uint64)(ALLOC_BEGIN + (KERN_PAGES * PGSIZE)); // 用户区域起始地址
    user_region.end = (uint64)(ALLOC_END-1); // 用户区域结束地址
    user_region.allocable = (user_region.end - user_region.begin) / PGSIZE +1; // 计算可分配页面数

    user_region.list_head.next = NULL; // 初始化链表头为 NULL
    spinlock_init(&user_region.lk, "user_region_lock"); // 初始化自旋锁

    //kernel的初始化
    init_range((void*)kern_region.begin,(void*)kern_region.end,true);
    //user的初始化
    init_range((void*)user_region.begin,(void*)user_region.end,false);
    
}

// 释放物理页
// 失败则panic锁死
static void pmem_free_first(uint64 page, bool in_kernel)
{
    alloc_region_t* ar = in_kernel ? &kern_region : &user_region;
    uint64 region_start = in_kernel ? (uint64)ALLOC_BEGIN : (uint64)(ALLOC_BEGIN + KERN_PAGES * PGSIZE);
    uint64 region_end = in_kernel ? (uint64)(ALLOC_BEGIN + KERN_PAGES * PGSIZE) : (uint64)ALLOC_END;

    // 检查页地址是否有效
    if ((page % PGSIZE) != 0 || (char*)page < (char*)region_start || page >= region_end) {
        panic("pmem_free: invalid page");
    }

    // 清空页面内容以便调试时捕获
    memset((void*)page, 1, PGSIZE);

    page_node_t p;
    p.next = (page_node_t*)page; // 将 page 转换为 page_node_t 指针

    p.next->refNum=0;

    // 上锁并插入到链表头
    spinlock_acquire(&ar->lk);
    p.next->next = ar->list_head.next;
    ar->list_head = p;
    spinlock_release(&ar->lk);
}

void init_range(void *pa_start, void *pa_end,bool in_kernel)
{
    char *ptr;
    ptr = (char *)ALIGN_DOWN((uint64)pa_end, PGSIZE);

    // 从大到小遍历，直到 ptr 小于 pa_start
    for (; ptr >= (char *)pa_start; ptr -= PGSIZE) {
        pmem_free_first((uint64)ptr, in_kernel);
    }
}

// 返回一个可分配的干净物理页
// 失败则panic锁死
//alloc操作就是移除 list_head->next
void* pmem_alloc(bool in_kernel)
{
    alloc_region_t* ar = in_kernel ? &kern_region : &user_region;
    page_node_t p;

    // 锁定区域
    spinlock_acquire(&ar->lk);

    // 获取链表头节点
    p.next = ar->list_head.next;

    if (p.next) {
        // 将链表头移动到下一个节点
        ar->list_head.next = p.next->next;
    }

    // 释放锁
    spinlock_release(&ar->lk);

    if (p.next) {
        // 如果成功获取到物理页，填充页面内容为 0 (调试用)
        memset((char*)p.next, 0, PGSIZE);
        // p.next->refNum=1;
    }

    // 返回分配的物理页地址（或 NULL）
    return (void*)p.next;
}

// 释放物理页
// 失败则panic锁死
void pmem_free(uint64 page, bool in_kernel)
{
    alloc_region_t* ar = in_kernel ? &kern_region : &user_region;
    uint64 region_start = in_kernel ? (uint64)ALLOC_BEGIN : (uint64)(ALLOC_BEGIN + KERN_PAGES * PGSIZE);
    uint64 region_end = in_kernel ? (uint64)(ALLOC_BEGIN + KERN_PAGES * PGSIZE) : (uint64)ALLOC_END;

    // 检查页地址是否有效
    if ((page % PGSIZE) != 0 || (char*)page < (char*)region_start || page >= region_end) {
        panic("pmem_free: invalid page");
    }

    // 清空页面内容以便调试时捕获
    memset((void*)page, 1, PGSIZE);

    page_node_t p;
    p.next = (page_node_t*)page; // 将 page 转换为 page_node_t 指针

    if (p.next->refNum>1)
    {
        p.next->refNum--;
        return;
    }

    // 上锁并插入到链表头
    spinlock_acquire(&ar->lk);
    p.next->next = ar->list_head.next;
    ar->list_head = p;
    spinlock_release(&ar->lk);
}


int cow_alloc(pgtbl_t pgtbl, uint64 va)
{
    va = ALIGN_DOWN(va,PGSIZE);
    pte_t* pte = vm_getpte(pgtbl, va, false);
    if(va >=VA_MAX ) return -1;
    if (pte == NULL || !(*pte & PTE_V)) {
        printf("cow_alloc: Invalid page table entry at va=0x%lx\n", va);
        return -1;
    }
    if (!(*pte & PTE_COW)) {
        printf("cow_alloc: Page is not marked as COW at va=0x%lx\n", va);
        return -1;
    }
    uint64 pa = PTE_TO_PA(*pte);
    page_node_t* page = (page_node_t*)pa;
    void* new_pa = pmem_alloc(false);
    if (new_pa == NULL) {
        printf("cow_alloc: pmem_alloc failed\n");
        return -1;
    }
     // 复制数据从旧页到新页
    memmove(new_pa, (char*)pa, PGSIZE);
    uint64 flags = PTE_FLAGS(*pte);
    flags = (flags | PTE_W) & ~PTE_COW;

    vm_mappages(pgtbl, va, (uint64)new_pa, PGSIZE, flags);
    page->refNum--;

    // 如果引用计数为零，释放原页
    if (page->refNum == 0) {
        pmem_free(pa,false);
    }
    return 0;
}